/*
 * Copyright(C) 2009, kesalin All rights reserved.
 * 
 * File name : DataObject.h
 * Brief     : Implement of IDropSource
 * Version   : 1.0
 * Author    : kesalin@gmail.com
 * Date      : 2009-6-5
 */

#pragma once

namespace DragDrop {

// DataItem
struct DataItem {

	FORMATETC Format;
	STGMEDIUM Medium;

}; // DataItem

// IDataObjectImpl
class IDataObjectImpl : public IDataObject {

	typedef ATL::CAtlList<DataItem> DataList;

	DataList m_Data;

	ATL::CComPtr<IDataAdviseHolder> m_DataAdviseHolder;

private:

	HRESULT Find(FORMATETC* fe, bool bAdd, DataItem** pp)
	{
		*pp = 0;

		if (fe->ptd != 0) {
			return DV_E_DVTARGETDEVICE;
		}

		POSITION pos = m_Data.GetHeadPosition();
		while (pos != 0) {
			DataItem& item = m_Data.GetNext(pos);

			if (item.Format.cfFormat == fe->cfFormat &&
				item.Format.dwAspect == fe->dwAspect &&
				item.Format.lindex   == fe->lindex) {
				if (!bAdd && (item.Format.tymed & fe->tymed) == 0) {
					return DV_E_TYMED;
				}

				CopyMemory(&item.Format, fe, sizeof(FORMATETC));

				*pp = &item;

				return S_OK;
			}
		}

		if (!bAdd) {
			return DV_E_FORMATETC;
		}

		DataItem& newItem = m_Data.GetAt(m_Data.AddTail());

		ZeroMemory(&newItem, sizeof(DataItem));
		CopyMemory(&newItem.Format, fe, sizeof(FORMATETC));

		*pp = &newItem;

		return S_OK;
	}

	static void ReleaseMedium(STGMEDIUM* m)
	{
		if (m != 0) {
			if (m->tymed != 0) {
				ReleaseStgMedium(m);
			}
			ZeroMemory(m, sizeof(STGMEDIUM));
		}
	}

	void ReleaseData()
	{
		POSITION pos = m_Data.GetHeadPosition();
		while (pos != 0) {
			DataItem& item = m_Data.GetNext(pos);
			CoTaskMemFree(item.Format.ptd);
			ReleaseMedium(&(item.Medium));
		}

		m_Data.RemoveAll();
	}

	/* */

	static HGLOBAL CopyGlobal(HGLOBAL src)
	{
		HGLOBAL dst = 0;

		void* ss = GlobalLock(src);
		if (ss != 0) {
			SIZE_T cb = GlobalSize(src);
			dst = GlobalAlloc(GPTR, cb);
			if (dst != 0 && cb > 0) {
				CopyMemory(dst, ss, cb);
			}
			GlobalUnlock(src);
		}

		return dst;
	}

	HRESULT CopyMedium(STGMEDIUM* src, STGMEDIUM* dst, bool bCopy)
	{
		STGMEDIUM dd = *src;

		if (src->pUnkForRelease == 0 &&
			(src->tymed & (TYMED_ISTREAM | TYMED_ISTORAGE)) == 0) {
			if (bCopy) {
				if (src->tymed != TYMED_HGLOBAL) {
					return DV_E_TYMED;
				}

				dd.hGlobal = CopyGlobal(src->hGlobal);
				if (dd.hGlobal == 0) {
					return E_OUTOFMEMORY;
				}

			} else {
				dd.pUnkForRelease = static_cast<IDataObject*>(this);
			}
		}

		switch (dd.tymed) {
		case TYMED_ISTREAM:
			dd.pstm->AddRef();
			break;

		case TYMED_ISTORAGE:
			dd.pstg->AddRef();
			break;
		}

		if (dd.pUnkForRelease != 0) {
			dd.pUnkForRelease->AddRef();
		}

		*dst = dd;

		return S_OK;
	}

	/* */

	static void* GetIdentity(IUnknown* p)
	{
		IUnknown* punk = 0;
		if (p != 0) {
			p->QueryInterface(IID_IUnknown, (void**)&punk);
			if (punk != 0) {
				punk->Release();
			}
		}
		return punk;
	}

	static bool IsEqualIdentity(IUnknown* p1, IUnknown* p2)
	{
		void* id1 = GetIdentity(p1);
		void* id2 = GetIdentity(p2);
		return (id1 == id2);
	}

	/* */

public:

	/* */

	IDataObjectImpl()
	{
	}

	~IDataObjectImpl()
	{
		ReleaseData();
	}

	/* */

	STDMETHOD(GetData)(
		FORMATETC* pformatetc,
		STGMEDIUM* pmedium)
	{
		if (pformatetc == 0) {
			return E_POINTER;
		}

		DataItem* item = 0;
		HRESULT hRslt = Find(pformatetc, false, &item);
		if (FAILED(hRslt)) {
			return hRslt;
		}

		if (pmedium == 0) {
			return E_POINTER;
		}

		hRslt = CopyMedium(&item->Medium, pmedium, false);
		if (FAILED(hRslt)) {
			return hRslt;
		}

		return S_OK;
	}

	/* */

	STDMETHOD(GetDataHere)(
		FORMATETC* /*pformatetc*/,
		STGMEDIUM* /*pmedium*/)
	{
		return E_NOTIMPL;
	}

	/* */

	STDMETHOD(QueryGetData)(
		FORMATETC* pformatetc)
	{
		if (pformatetc == 0) {
			return E_POINTER;
		}

		DataItem* item = 0;
		HRESULT hRslt = Find(pformatetc, false, &item);
		if (FAILED(hRslt)) {
			return hRslt;
		}

		return S_OK;
	}

	/* */

	STDMETHOD(GetCanonicalFormatEtc)(
		FORMATETC* /*pformatectIn*/,
		FORMATETC* /*pformatetcOut*/)
	{
		return E_NOTIMPL;
	}

	/* */

	STDMETHOD(SetData)(
		FORMATETC* pformatetc,
		STGMEDIUM* pmedium,
		BOOL       fRelease)
	{
		if (pformatetc == 0) {
			return E_POINTER;
		}

		DataItem* item = 0;
		HRESULT hRslt = Find(pformatetc, true, &item);
		if (FAILED(hRslt)) {
			return hRslt;
		}

		if (pmedium == 0) {
			return E_POINTER;
		}

		ReleaseMedium(&item->Medium);

		if (fRelease) {
			item->Medium = *pmedium;

		} else {
			hRslt = CopyMedium(pmedium, &item->Medium, true);
			if (FAILED(hRslt)) {
				return hRslt;
			}
		}

		item->Format.tymed = item->Medium.tymed;

		if (IsEqualIdentity(
			item->Medium.pUnkForRelease,
			static_cast<IDataObject*>(this))) {
			item->Medium.pUnkForRelease->Release();
			item->Medium.pUnkForRelease = 0;
		}

		return S_OK;
	}

	/* */

	STDMETHOD(EnumFormatEtc)(
		DWORD            dwDirection,
		IEnumFORMATETC** ppenumFormatEtc)
	{
		if (ppenumFormatEtc == 0) {
			return E_POINTER;
		}

		*ppenumFormatEtc = 0;

		if (dwDirection != DATADIR_GET) {
			return E_NOTIMPL;
		}

		SIZE_T count = m_Data.GetCount();
		if (count == 0) {
			return E_UNEXPECTED;
		}

		ATL::CAtlArray<FORMATETC> fmts;
		fmts.SetCount(count);

		FORMATETC* pfmts = fmts.GetData();

		POSITION pos = m_Data.GetHeadPosition();
		while (pos != 0) {
			const DataItem& item = m_Data.GetNext(pos);
			*(pfmts++) = item.Format;
		}

		HRESULT hRslt = CreateFormatEnumerator(
			UINT(count),
			fmts.GetData(),
			ppenumFormatEtc);
		if (FAILED(hRslt)) {
			return hRslt;
		}

		return S_OK;
	}

	/* */

	STDMETHOD(DAdvise)(
		FORMATETC*   pformatetc,
		DWORD        advf,
		IAdviseSink* pAdvSink,
		DWORD*       pdwConnection)
	{
		HRESULT hRslt = S_OK;

		if (m_DataAdviseHolder == 0) {
			hRslt = CreateDataAdviseHolder(&m_DataAdviseHolder);
			if (hRslt != S_OK) {
				return hRslt;
			}
		}

		hRslt = m_DataAdviseHolder->Advise(
			static_cast<IDataObject*>(this),
			pformatetc,
			advf,
			pAdvSink,
			pdwConnection);
		if (FAILED(hRslt)) {
			return hRslt;
		}

		return S_OK;
	}

	STDMETHOD(DUnadvise)(
		DWORD dwConnection)
	{
		if (m_DataAdviseHolder == 0) {
			return OLE_E_NOCONNECTION;
		}

		HRESULT hRslt = m_DataAdviseHolder->Unadvise(dwConnection);
		if (FAILED(hRslt)) {
			return hRslt;
		}

		return S_OK;
	}

	STDMETHOD(EnumDAdvise)(
		IEnumSTATDATA** ppenumAdvise)
	{
		if (ppenumAdvise == 0) {
			return E_POINTER;
		}

		*ppenumAdvise = 0;

		if (m_DataAdviseHolder == 0) {
			return E_FAIL;
		}

		HRESULT hRslt = m_DataAdviseHolder->EnumAdvise(ppenumAdvise);
		if (FAILED(hRslt)) {
			return hRslt;
		}

		return S_OK;
	}

	/* */

}; // IDataObjectImpl

} // namespace DragDrop

